import cv2
import numpy as np


def find_image_position(template_path, screenshot_path):
    """
    使用特征点匹配找到模板图片在屏幕截图中的位置
    
    Args:
        template_path (str): 模板图片路径 (at.png)
        screenshot_path (str): 屏幕截图路径 (3.png)
    
    Returns:
        tuple: 包含匹配位置的坐标 (top_left, bottom_right) 或 None（如果没有找到）
    """
    
    # 读取模板图片和屏幕截图
    template = cv2.imread(template_path, cv2.IMREAD_GRAYSCALE)
    screenshot = cv2.imread(screenshot_path, cv2.IMREAD_GRAYSCALE)
    
    if template is None:
        raise FileNotFoundError(f"无法读取模板图片: {template_path}")
    
    if screenshot is None:
        raise FileNotFoundError(f"无法读取屏幕截图: {screenshot_path}")
    
    # 创建SIFT特征检测器
    sift = cv2.SIFT_create()
    
    # 检测关键点并计算描述符
    kp1, des1 = sift.detectAndCompute(template, None)
    kp2, des2 = sift.detectAndCompute(screenshot, None)
    
    if des1 is None or des2 is None:
        print("未能检测到足够的特征点")
        return None
    
    # 使用FLANN匹配器进行特征匹配
    FLANN_INDEX_KDTREE = 1
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    search_params = dict(checks=50)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    
    matches = flann.knnMatch(des1, des2, k=2)
    
    # 应用 Lowe's ratio test 来过滤好的匹配点
    good_matches = []
    for pair in matches:
        if len(pair) == 2:
            m, n = pair
            if m.distance < 0.7 * n.distance:
                good_matches.append(m)
    
    # 如果有足够的匹配点，则计算位置
    if len(good_matches) > 10:
        # 提取匹配点的坐标
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        
        # 计算单应性矩阵
        M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
        
        if M is not None:
            # 获取模板图片的尺寸
            h, w = template.shape
            
            # 定义模板图片的四个角点
            pts = np.float32([[0, 0], [0, h-1], [w-1, h-1], [w-1, 0]]).reshape(-1, 1, 2)
            
            # 使用单应性矩阵透视变换角点
            dst = cv2.perspectiveTransform(pts, M)
            
            # 计算边界框坐标
            min_x = int(np.min(dst[:, 0, 0]))
            max_x = int(np.max(dst[:, 0, 0]))
            min_y = int(np.min(dst[:, 0, 1]))
            max_y = int(np.max(dst[:, 0, 1]))
            
            top_left = (min_x, min_y)
            bottom_right = (max_x, max_y)
            
            return top_left, bottom_right
    
    return None


def draw_rectangle_on_screenshot(screenshot_path, position, output_path):
    """
    在屏幕截图上绘制矩形标记匹配位置
    
    Args:
        screenshot_path (str): 屏幕截图路径
        position (tuple): 包含(top_left, bottom_right)坐标的元组
        output_path (str): 输出图片路径
    """
    screenshot = cv2.imread(screenshot_path)
    top_left, bottom_right = position
    cv2.rectangle(screenshot, top_left, bottom_right, (0, 255, 0), 2)
    cv2.imwrite(output_path, screenshot)


if __name__ == "__main__":
    # 定义文件路径
    template_img_path = r"D:\PycharmProjects\auto-script\template\APP-Open.png"
    screenshot_path = r"D:\PycharmProjects\auto-script\screentshot\APP-Open.png"
    output_path = r"D:\PycharmProjects\auto-script\screentshot\APP-Open_marked.png"
    
    try:
        # 查找图片位置
        position = find_image_position(template_img_path, screenshot_path)
        
        if position:
            top_left, bottom_right = position
            print(f"找到匹配位置:")
            print(f"左上角坐标: {top_left}")
            print(f"右下角坐标: {bottom_right}")
            
            # 在截图上绘制矩形并保存
            draw_rectangle_on_screenshot(screenshot_path, position, output_path)
            print(f"标记结果已保存到: {output_path}")
        else:
            print("未找到匹配位置")
            
    except Exception as e:
        print(f"发生错误: {e}")